package org.theSeed.webFlux.filter;

import cn.hutool.core.util.ObjectUtil;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import org.theSeed.base.pojo.constant.SysConstant;
import org.theSeed.webFlux.sysEnum.FilterOrder;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.List;

/**
 * 可重复读过滤器
 */
@Slf4j
public class RepeatReadFilter implements WebFilter, Ordered {
    public RepeatReadFilter() {
    }

    public RepeatReadFilter(Integer orderNum) {
        this.orderNum = orderNum;
    }

    protected Integer orderNum = FilterOrder.REPEAT_READ.getCode();

    @Override
    public int getOrder() {
        return orderNum;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        ServerHttpResponse response = exchange.getResponse();
        MediaType contentType = request.getHeaders().getContentType();
        //请求体转换
        Mono<ServerHttpRequest> requestMono = null;
        if (CachedBodyHttpRequestWrapper.isRequestBodyCompatible(request)){
            requestMono = request.getBody()
                    .collectList()
                    .map(dataBuffers -> {
                        CachedBodyHttpRequestWrapper cachedBodyHttpRequestWrapper = new CachedBodyHttpRequestWrapper(request, dataBuffers);
                        return cachedBodyHttpRequestWrapper;
                    });
        }else
            requestMono = Mono.just(request);
        //响应体转换
        CachedBodyHttpResponseWrapper cachedBodyHttpResponseWrapper = new CachedBodyHttpResponseWrapper(response);
        return requestMono.flatMap(req->{
            return chain.filter(exchange.mutate().request(req).response(cachedBodyHttpResponseWrapper).build());
        });
    }

    /*
     * 可重复读取请求体的 HttpRequest 包装器
    */
    public static class CachedBodyHttpRequestWrapper extends ServerHttpRequestDecorator {
        private final List<DataBuffer> cachedBody;

        public CachedBodyHttpRequestWrapper(ServerHttpRequest delegate, List<DataBuffer> cachedBody) {
            super(delegate);
            this.cachedBody = cachedBody;
        }

        @Override
        public Flux<DataBuffer> getBody() {
            // 返回缓存的请求体，可以多次读取
            return Flux.fromIterable(cachedBody);
        }

        public Mono<DataBuffer> getCachedBody(){
            return DataBufferUtils.join(Flux.fromIterable(cachedBody));
        }

        /**
         * 判断请求是否具有可读取的请求体
         * 可以根据实际需求调整条件
         */
        public static boolean isRequestBodyCompatible(ServerHttpRequest serverHttpRequest) {
            MediaType contentType = serverHttpRequest.getHeaders().getContentType();
            if (ObjectUtil.isNotEmpty(contentType)){
                HttpMethod method = serverHttpRequest.getMethod();
                // 通常只有 POST, PUT, PATCH 等方法才有请求体
                // 并且 content-type 不为空（例如排除 multipart/form-data，因为文件上传处理方式不同）
                return ((HttpMethod.POST.equals(method) || HttpMethod.PUT.equals(method) || HttpMethod.DELETE.equals(method) && (MediaType.APPLICATION_JSON.equalsTypeAndSubtype(contentType) || MediaType.APPLICATION_XML_VALUE.equals(contentType))))
                        || HttpMethod.GET.equals(method);
            }else
                return true;
        }
    }

    /**
     * 可缓存响应体的 HttpResponse 包装器
     */
    public static class CachedBodyHttpResponseWrapper extends ServerHttpResponseDecorator {

        private final StringBuilder bodyCache = new StringBuilder();
        private final List<DataBuffer> dataBuffers = new ArrayList<>();
        private Boolean responseBodyCompatible = null;

        public CachedBodyHttpResponseWrapper(ServerHttpResponse delegate) {
            super(delegate);
        }

        @Override
        public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
            if (body instanceof Flux) {
                Flux<DataBuffer> fluxBody = (Flux<DataBuffer>) body;
                // 缓存响应体
                Flux<DataBuffer> cachedFlux = fluxBody;
                if (isResponseBodyCompatible(this)){
                    cachedFlux.doOnNext(buffer -> {
                        if (responseBodyCompatible) {
                            cacheDataBuffer(buffer);
                        }
                    })
                    .cache(); // 使响应体可重复读取
                }
                return super.writeWith(cachedFlux);
            }
            return super.writeWith(body);
        }

        @Override
        public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
            // 处理分块响应
            return super.writeAndFlushWith(body);
        }

        private void cacheDataBuffer(DataBuffer buffer) {
            try {
                // 读取缓冲区内容
                byte[] bytes = new byte[buffer.readableByteCount()];
                buffer.read(bytes);
                String content = new String(bytes, SysConstant.defaultCharset);

                // 添加到缓存
                bodyCache.append(content);
                dataBuffers.add(buffer);

            } catch (Exception e) {
                log.warn("Failed to cache response buffer", e);
            }
        }

        public String getCachedBody() {
            return bodyCache.toString();
        }

        public List<DataBuffer> getCachedDataBuffers() {
            return new ArrayList<>(dataBuffers);
        }

        public boolean isContentCached() {
            return bodyCache.length() > 0;
        }

        public void disableCaching() {
            this.responseBodyCompatible = false;
        }

        public void clearCache() {
            bodyCache.setLength(0);
            dataBuffers.forEach(DataBufferUtils::release);
            dataBuffers.clear();
        }

        /**
         * 排除某些不需要处理的请求（例如静态资源、文件下载等）
         */
        protected boolean isResponseBodyCompatible(ServerHttpResponse serverHttpResponse){
            MediaType contentType = serverHttpResponse.getHeaders().getContentType();
            if (ObjectUtil.isNotNull(contentType)){
                return MediaType.APPLICATION_JSON.equalsTypeAndSubtype(contentType) || MediaType.APPLICATION_XML.equalsTypeAndSubtype(contentType);
            }else{
                return false;
            }
        }
    }
}
